clear;
clc;
%% Preliminary
Folder      =   ['TableGraphs\\IRF_Decomposition\\']; mkdir(Folder);
FileName    =   'Result_Baseline';
StdVec      =   [1 1 1 -1 -1 60]'*0.01/4;
% PP          =   Setup_PP(struct('KAPPA',0.00/4,'XI',1/82.5/4,'Pir_SS',0.02/4,...
%                                 'Flag_dE',1,'Flag_Price',0,...
%                                 'B_Total_SS',0.95,'b_lb',-0.21,...
%                                 'Prob_H',0.5334,...
%                                 'RHO_M_dom',0.68,'RHO_M_ext',0.81,'RHO_Y_H',0.44,...
%                                 'Taylor_Pi',1.1,'Taylor_ir',0.87,...
%                                 'Cons_ElasTN',4.62,'Cons_ElasHF',4.62));
PP          =   Setup_PP(struct('KAPPA',0.00/4,'XI',1/82.5/4,'Pir_SS',0.02/4,...
                                'Flag_dE',1,'Flag_Price',0,...
                                'B_Total_SS',0.91,'b_lb',-0.29,...
                                'TrProb_H',0.99,'Prob_H',0.35,...
                                'TrProb_ext',0.99,'Prob_ext',0.16,...
                                'Cons_FracT',0.3248,'Cons_FracH',0.60,...
                                'RHO_M_dom',0.68,'RHO_M_ext',0.81,'RHO_Y_H',0.95,...
                                'Taylor_Pi',1.1,'Taylor_ir',0.87,...
                                'Cons_ElasTN',6.19,'Cons_ElasHF',6.19));
                            
                            
SOLUTION    =   Solver_Main(PP,StdVec);

Num_T       =   40;
%% Decompose the IRF
VarList     =   {'Pir','dE','T','Profit','TAU','ir_dom','ir_ext','w'};
ShockList   =   {'Eps_Eta_M_dom','Eps_Eta_M_ext','Eps_Eta_Y_H'};

InputDev    =   struct();
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    for jj=1:length(VarList)
        VV          =   VarList{jj};
        InputDev.(SH).(VV)  =   SOLUTION.IRF.(SH).(VV)(:,1:Num_T);
    end
end

DecDev      =   struct();
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    DecDev.(SH) =   SubFun_UnitIRFDec(SOLUTION.PP,SOLUTION.SS,SOLUTION.MODEL,VarList,InputDev.(SH));
end

%% Group-wise Decomposition
% Setup
VarGroup    =   struct('Interest_Domestic',{{'ir_dom'}},...
                       'Interest_Foreign',{{'ir_ext'}},...
                       'Inflation',{{'Pir'}},...
                       'ExchangeRate',{{'dE'}},'w',{{'w'}},'TAU',{{'TAU'}},...
                       'T',{{'T'}},'Profit',{{'Profit'}} ...
                       );
VarGroupLabel=  struct('Interest_Domestic','Interest Rate, Domestic',...
                       'Interest_Foreign','Interest Rate, Foreign',...
                       'Inflation','Inflation',...
                       'ExchangeRate','Exchange Rate','w','Wage','TAU','Tax',...
                       'T','Transfer','Profit','Dividend' ...
                       );
GroupList   =   fields(VarGroup);
GroupLabelList= cell(length(GroupList),1);
for ii=1:length(GroupList)
    GroupLabelList{ii}  =   VarGroupLabel.(GroupList{ii});
end
GroupDecDev =   struct();
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    TempList    =   {'B','C','L'};
    for jj=1:length(TempList)
        VV          =   TempList{jj};
        GroupDecDev.(SH).(VV) ...
                    =   zeros(SOLUTION.MODEL.VAR.(VV).Dim,Num_T,length(GroupList));
        for gg=1:length(GroupList)
            GG          =   GroupList{gg};
            GGTempList  =   VarGroup.(GG);
            for kk=1:length(GGTempList)
                GroupDecDev.(SH).(VV)(:,:,gg)    =   GroupDecDev.(SH).(VV)(:,:,gg)+DecDev.(SH).(GGTempList{kk}).(VV);
            end
        end
    end
end

%% Collect the Results
TempVarList =   {'C','B','L'};
IrfAgg      =   struct();
IrfDec      =   struct();
IrfDecTable =   struct();
Table_Rows  =   GroupList;
Table_Cols  =   {'C';'B';'L'};
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    % Peak IRF of Aggregate C
    [~,AggInd]  =   max(abs(SOLUTION.IRF.(SH).C(:,1:Num_T)));
    % First 8 Quarters
%     AggInd      =   (1:8);
    % Normalizer
    Normalizer  =   abs(sum(SOLUTION.IRF.(SH).C(:,AggInd)));
    for jj=1:length(TempVarList)
        vv          =   TempVarList{jj};
        TempSize    =   size(GroupDecDev.(SH).(vv));
        IrfAgg.(SH).(vv) ...
                    =   reshape(sum(GroupDecDev.(SH).(vv)(:,AggInd,:),2),TempSize([1,3]));
        IrfDec.(SH).(vv) ...
                    =   IrfAgg.(SH).(vv)./sum(IrfAgg.(SH).(vv),2);
    end
    %----------------------------------------------------------------------
    % Tables
    %----------------------------------------------------------------------
    % Aggregated
    Table_Data  =   [IrfAgg.(SH).C;IrfAgg.(SH).B;IrfAgg.(SH).L]';
    IrfAggTable.(SH) ...
                =   array2table(Table_Data,'VariableNames',Table_Cols,'RowNames',Table_Rows);
    % Pct
    IrfDecTable.(SH) ...
                =   IrfAggTable.(SH);
    for jj=1:length(IrfAggTable.(SH).Properties.VariableNames)
        vv          =   IrfAggTable.(SH).Properties.VariableNames{jj};
        IrfDecTable.(SH).(vv) ...
                    =   IrfDecTable.(SH).(vv)/sum(IrfDecTable.(SH).(vv));
    end
    % Print Tables
     for jj=1:length(IrfAggTable.(SH).Properties.VariableNames)
        vv          =   IrfAggTable.(SH).Properties.VariableNames{jj};
        IrfAggTable.(SH).(vv)   =   round(IrfAggTable.(SH).(vv),6);
        IrfDecTable.(SH).(vv)   =   round(IrfDecTable.(SH).(vv),3);
    end
    writetable(IrfAggTable.(SH),[Folder,'IrfAgg.xlsx'],'WriteRowNames',true,'Sheet',SH);
    writetable(IrfDecTable.(SH),[Folder,'IrfDec.xlsx'],'WriteRowNames',true,'Sheet',SH);
end


%% Visual Decomposition 
N_col       =   1;
N_row       =   1;
fig_size    =   [N_col,N_row].*[0.75,1]*1/2; 
fig_gap     =   [0.05,0.08];
fig_VMargin =   [0.05,0.03];
fig_HMargin =   [0.10,0.02]; 

fig_SubPlot =   @(ii_vv)subtightplot(N_row,N_col,ii_vv,...
                                     fig_gap,fig_VMargin,fig_HMargin);
AxFontSize  =   8;
TickFmt     =   '%.2f';
XLabel      =   @(Label)xlabel(Label,'fontsize',8,'interpreter','latex');
YLabel      =   @(Label)ylabel(Label,'fontsize',8,'interpreter','latex');
LEG         =   @(PH,LabelList)legend(PH,LabelList,'fontsize',8,'interpreter','latex','location','northeast');


for ii=1:length(ShockList)
    SH = ShockList{ii};
    
    YNames = cell(length(IrfAggTable.(SH).Row),1);
    for jj=1:length(YNames); YNames{jj} = VarGroupLabel.(IrfAggTable.(SH).Row{jj}); end
    % Financial Integration
    TempData = [IrfAggTable.(SH).C,IrfAggTable.(SH).L,IrfAggTable.(SH).B];
    XNames = {'Consumption','Hours','Saving'};

    Fig = FigureSetup([SH,', IrfDec'],fig_size);
    ax = fig_SubPlot(1);
    [~,ax] = Plot_StackedBars(TempData,YNames,XNames,[],ax);
    
    print('-depsc','-r1000',Fig,[Folder,'IrfDec_',SH]);
    close(Fig);

end

%% Decomposed IRFs
N_col       =   1;
N_row       =   1;
fig_size    =   [N_col,N_row].*[0.75,1]*1/2; 
fig_gap     =   [0.05,0.08];
fig_VMargin =   [0.08,0.03];
fig_HMargin =   [0.15,0.02]; 

fig_SubPlot =   @(ii_vv)subtightplot(N_row,N_col,ii_vv,...
                                     fig_gap,fig_VMargin,fig_HMargin);
AxFontSize  =   8;
TickFmt     =   '%.2f';
XLabel      =   @(Label)xlabel(Label,'fontsize',8,'interpreter','latex');
YLabel      =   @(Label)ylabel(Label,'fontsize',8,'interpreter','latex');
LEG         =   @(PH,LabelList)legend(PH,LabelList,'fontsize',8,'interpreter','latex','location','northeast');


VarList     =   {'C','L','B'};
LineStyleList   =   {'-','--','--','--','--',':',':',':',':';...
                     'black', 'blue','red','m','c','g','y','b','r';...
                     3,2,2,2,2,2,2,2,2};
LabelList   =   {'Aggregate', ...
                 'Interest Rate, Domestic','Interest Rate, Foreign',...
                 'Inflation','Exchange Rate','Wage','Tax Rate','Transfer','Dividend'};
for ii=1:length(ShockList)
    SH          =   ShockList{ii};
    for jj=1:length(VarList)
        VV      =   VarList{jj};
        
        TempDec     =   [DecDev.(SH).ir_dom.(VV);...
                         DecDev.(SH).ir_ext.(VV);...
                         DecDev.(SH).Pir.(VV);...
                         DecDev.(SH).dE.(VV);...
                         DecDev.(SH).w.(VV);...
                         DecDev.(SH).TAU.(VV);...
                         DecDev.(SH).T.(VV);...
                         DecDev.(SH).Profit.(VV)];
        TempAgg     =   SOLUTION.IRF.(SH).(VV)(1,1:Num_T);
        
        Fig         =   FigureSetup(['IRF: ',VV,' to ',SH],fig_size);
        ax          =   fig_SubPlot(1);
        TempIRF     =   [TempAgg;TempDec]*100;
        [PH,ax]     =   IRFSoloPlot(TempIRF(:,1:12),LineStyleList,0,ax);
        XLabel('Quarter'); YLabel('Deviation from Steady State');
        ax.YAxis.Exponent   =   0;
        ax.YAxis.TickLabelFormat    =   TickFmt;
        ax.FontSize =   AxFontSize;
        LEG(PH,LabelList);
        print('-depsc','-r1000',Fig,[Folder,'IRF_',VV,'_',SH]);
        close(Fig);
    end
end


